import torch
import torch.nn as nn
from torch.nn import functional as F

F.one_hot()


class CasModel(nn.Module):
    def __init__(self, num_nodes, embedding_dim, hidden_size, num_layers):
        super().__init__()
        self.emb = nn.Embedding(num_nodes, embedding_dim, -1)
        self.gru = nn.GRU(embedding_dim, hidden_size, num_layers, batch_first=True, bidirectional=True)
        self.p_geo = nn.Parameter(torch.tensor(0, requires_grad=True), requires_grad=True)
        self.att_lambda = nn.Parameter(torch.randn())

    def forward(self, x, sz):
        pass


class NodeEmbedding(nn.Module):
    def __init__(self, num_classes, embedding_length):
        super().__init__()
        self.num_classes = num_classes
        self.emb = nn.Linear(num_classes, embedding_length)

    def forward(self, x):
        inputs = F.one_hot(x, self.num_classes)
        outputs = self.emb(inputs)
        return outputs


def run(args):
    return None
